#from libmyatomism import *
#from math import *
import array 
import matplotlib.pyplot as plt
from libatomism import *
import math

msLogger.setPriorityMin(Priority.ERROR)

unit=msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

q0 = msScalarVariable.New("Angstrom").set(5,0.5,6,0.1,0.1).setId("q0")
q1 = msScalarVariable.New("Degree").set(120,0,360,2,2).setId("q1")
q2 = msScalarVariable.New("Degree").set(120,0,360,2,2).setId("q2")
q3 = msScalarVariable.New("Degree").set(120,0,360,2,2).setId("q3")

gcoors = msGeneralizedCoordinates.New(unit).addVar(q0).addVar(q1).addVar(q2).addVar(q3)

zmat = msZMat.New(unit).setId("zmat").set("N N 0 1.1 N 1 5. 0 10. N 2 1.1 1 90. 0 0.")
system=msSystem.New(unit).addEntity(zmat)
kinfct0 = msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0")
kinfct1 = msParser.New("Degree").setCoordinates(gcoors).setExpression("q1").setId("f=q1")
kinfct2 = msParser.New("Degree").setCoordinates(gcoors).setExpression("q2").setId("f=q2")
kinfct3 = msParser.New("Degree").setCoordinates(gcoors).setExpression("q3").setId("f=q3")

kinop=msKineticOperator.New(unit).set(gcoors,system)
kinop.addDynamicDof(zmat,1,kinfct0).addDynamicDof(zmat,2,kinfct1).addDynamicDof(zmat,4,kinfct2).addDynamicDof(zmat,5,kinfct3)

#
#  Read file
#
with open('datafile.txt') as f:
    w = f.readline().split() # read first line
    array = []
    for line in f: # read rest of lines
        array.append([float(x) for x in line.split()])

repro =msRepromodel.New("")
repro.setCoordinates(gcoors)
d=0.25
N=4
Emax=100
#
#  create database
#
q0g = msScalarVariable.New("Angstrom").set(5,0.5,10,0.1,0.1).setId("q0g")
q1g = msScalarVariable.New("Angstrom").set(5,0.5,10,0.1,0.1).setId("q1g")
q2g = msScalarVariable.New("Angstrom").set(5,0.5,10,0.1,0.1).setId("q2g")
q3g = msScalarVariable.New("Angstrom").set(5,0.5,10,0.1,0.1).setId("q3g")
gcoorsg = msGeneralizedCoordinates.New(unit).addVar(q0g).addVar(q1g).addVar(q2g).addVar(q3g)

epot= msParser.New("kcal/mol").setCoordinates(gcoorsg)
epot.setExpression("4*e0*(  (sigma/q0g)^a-(sigma/q0g)^b + (sigma/q1g)^a-(sigma/q1g)^b + (sigma/q2g)^a-(sigma/q2g)^b + (sigma/q3g)^a-(sigma/q3g)^b )")
#epot.setExpression("0")
epot.setConstant("e0",0.2).setConstant("sigma",3.2).setConstant("a",9.2).setConstant("b",6)
 
for values in array:
    vec = VectorOfDouble()
    vec.extend( arg for arg in values)
    gcoors.setCoordinates(vec)
    kinop.setDynamicDofs()
    system.computeCartCoordinates()
    vec = VectorOfDouble()
    vec.extend([math.exp(-zmat.separation(0,2)/d),math.exp(-zmat.separation(0,3)/d),math.exp(-zmat.separation(1,2)/d),math.exp(-zmat.separation(1,3)/d)])
    vec2 = VectorOfDouble()
    vec2.extend([zmat.separation(0,2),zmat.separation(0,3),zmat.separation(1,2),zmat.separation(1,3)])
            
    if( values[5]<Emax ):
        repro.addData( vec, values[5] - epot.evaluate(vec2) )
 
#
#  fit
#
msLogger.setPriorityMin(Priority.INFO)
msLogger.setVerboseLevel(1)
repro.setBasisSet(N)
msLogger.setPriorityMin(Priority.ERROR)

repro.fit()
sampler=repro.getSampler()
sampler.begin()
field=repro.getField()
y=[]
x=[]
n=0
sigma=0
while(sampler.getNextPoint(gcoors)):
     x.append(field.evaluate())
     y.append(repro.evaluate())
     sigma+=abs(field.evaluate() - repro.evaluate())
     n=n+1
print "sigma=",sigma/n
plt.plot(x,y,"+")
plt.show()
     
    
#
#  validation
#
zmat = msZMat.New(unit).setId("zmat").set("X N 0 0.55 N 0 0.55 1 179.9 X 0 5. 1 90. 2 0. N 3 0.55 0 90. 1 180. N 3 0.55 4 179.9 0 180.")
system=msSystem.New(unit).addEntity(zmat)

kinop=msKineticOperator.New(unit).set(gcoors,system)
kinop.addDynamicDof(zmat,3,kinfct0).addDynamicDof(zmat,4,kinfct1).addDynamicDof(zmat,7,kinfct2).addDynamicDof(zmat,8,kinfct3)

scenarios=[[0.1,0.1,0],[90.,0.1,0],[90,90,0],[90,90,90]]
indice=[[0,0],[0,1],[1,0],[1,1]]
fig, axes = plt.subplots(2, 2)
i=0
for scenar in scenarios:
    filename="approach_"+str(scenar[0])+"_"+str(scenar[1])+"_"+str(scenar[2])+".txt"
    file = open(filename,'r')
    array = []
    with open(filename) as f:
         w = f.readline().split()
         for line in f: 
             array.append([float(x) for x in line.split()])
    y=[]
    data=[]
    x=[]
    for values in array:
         vec = VectorOfDouble()
         vec.extend( arg for arg in values)
         gcoors.setCoordinates(vec)
         kinop.setDynamicDofs()
         system.computeCartCoordinates()
    
         vec2 = VectorOfDouble()
         vec2.extend([math.exp(-zmat.separation(1,4)/d),math.exp(-zmat.separation(1,5)/d),math.exp(-zmat.separation(2,4)/d),math.exp(-zmat.separation(2,5)/d)])
         vec3 = VectorOfDouble()
         vec3.extend([zmat.separation(1,4),zmat.separation(1,5),zmat.separation(2,4),zmat.separation(2,5)])

         y.append( repro.evaluate(vec2) + epot.evaluate(vec3))
         x.append(vec[0])
         data.append(vec[5])
    axes[indice[i][0], indice[i][1]].plot(x,y,"-",x,data,"*r") 
    axes[indice[i][0],indice[i][1]].axis((0,6.,-1,Emax))
    i=i+1

fig.show()    